{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# EKS CSI FSX Lustre Setup\n",
    "\n",
    "Amazon FSx for Lustre is a high-performance file system optimized for deep learning workloads. FSx provides POSIX-compliant file system access to S3 for multiple readers and writers simultaneously.\n",
    "  \n",
    "The Amazon FSx for Lustre Container Storage Interface (CSI) driver provides a CSI interface that allows Amazon EKS clusters to manage the lifecycle of Amazon FSx for Lustre file systems.  \n",
    "\n",
    "* https://docs.aws.amazon.com/eks/latest/userguide/fsx-csi.html\n",
    "* https://github.com/kubernetes-sigs/aws-fsx-csi-driver"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import boto3\n",
    "import json\n",
    "from botocore.exceptions import ClientError\n",
    "\n",
    "iam = boto3.client(\"iam\")\n",
    "sts = boto3.client(\"sts\")\n",
    "cfn = boto3.client(\"cloudformation\")\n",
    "eks = boto3.client(\"eks\")\n",
    "\n",
    "region = boto3.Session().region_name\n",
    "cluster_name = \"workshop\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 1. Install the FSx CSI Driver for Kubernetes"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create IAM Policy\n",
    "\n",
    "Create an IAM policy and service account that allows the driver to make calls to AWS APIs on your behalf."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "!pygmentize fsx/fsx-csi-driver.json"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# !aws iam create-policy \\\n",
    "#     --policy-name Amazon_FSx_Lustre_CSI_Driver \\\n",
    "#     --policy-document file://fsx/fsx-csi-driver.json"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"fsx/fsx-csi-driver.json\") as json_file:\n",
    "    data = json.load(json_file)\n",
    "    policy = json.dumps(data)\n",
    "\n",
    "try:\n",
    "    response = iam.create_policy(PolicyName=\"Amazon_FSx_Lustre_CSI_Driver\", PolicyDocument=policy)\n",
    "    print(\"[OK] Policy created.\")\n",
    "\n",
    "except ClientError as e:\n",
    "    if e.response[\"Error\"][\"Code\"] == \"EntityAlreadyExists\":\n",
    "        print(\"[OK] Policy already exists.\")\n",
    "    else:\n",
    "        print(\"Error: %s\" % e)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "account_id = sts.get_caller_identity()[\"Account\"]\n",
    "csi_policy_arn = \"arn:aws:iam::{}:policy/Amazon_FSx_Lustre_CSI_Driver\".format(account_id)\n",
    "print(csi_policy_arn)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create Kubernetes IAM Service Account\n",
    "\n",
    "Create a Kubernetes service account for the driver and attach the policy to the service account. Replacing the ARN of the policy with the ARN returned in the previous step."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## _The next cell runs for about 10min. Please be patient._"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "!eksctl create iamserviceaccount \\\n",
    "     --region $region \\\n",
    "     --name fsx-csi-controller-sa \\\n",
    "     --namespace kube-system \\\n",
    "     --cluster $cluster_name \\\n",
    "     --attach-policy-arn $csi_policy_arn \\\n",
    "     --approve"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cf_stack_name = \"eksctl-{}-addon-iamserviceaccount-kube-system-fsx-csi-controller-sa\".format(cluster_name)\n",
    "print(cf_stack_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "response = cfn.list_stack_resources(StackName=cf_stack_name)\n",
    "print(response)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "iam_role_name = response[\"StackResourceSummaries\"][0][\"PhysicalResourceId\"]\n",
    "print(iam_role_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "iam_role_arn = iam.get_role(RoleName=iam_role_name)[\"Role\"][\"Arn\"]\n",
    "print(iam_role_arn)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Deploy CSI Driver"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!kubectl apply -k \"github.com/kubernetes-sigs/aws-fsx-csi-driver/deploy/kubernetes/overlays/stable/?ref=master\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Patch the driver deployment to add the service account that you just created, replacing the ARN with the correct role ARN."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!kubectl annotate serviceaccount -n kube-system fsx-csi-controller-sa \\\n",
    " eks.amazonaws.com/role-arn=$iam_role_arn --overwrite=true"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Check S3 Bucket For FSX"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bucket = \"s3://fsx-container-demo\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!aws s3 mb $bucket"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!aws s3 ls $bucket"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!aws s3 ls $bucket --recursive"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Download Storage Class Manifest"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!curl -o storageclass.yaml https://raw.githubusercontent.com/kubernetes-sigs/aws-fsx-csi-driver/master/examples/kubernetes/dynamic_provisioning_s3/specs/storageclass.yaml"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Get VPC ID and Subnet ID"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash\n",
    "\n",
    "source ~/.bash_profile\n",
    "\n",
    "#### Get VPC ID\n",
    "export VPC_ID=$(aws ec2 describe-vpcs --filters \"Name=tag:Name,Values=eksctl-${AWS_CLUSTER_NAME}-cluster/VPC\" --query \"Vpcs[0].VpcId\" --output text)\n",
    "echo \"export VPC_ID=${VPC_ID}\" | tee -a ~/.bash_profile\n",
    "\n",
    "#### Get Subnet ID\n",
    "export SUBNET_ID=$(aws ec2 describe-subnets --filters \"Name=vpc-id,Values=${VPC_ID}\" --query \"Subnets[0].SubnetId\" --output text)\n",
    "echo \"export SUBNET_ID=${SUBNET_ID}\" | tee -a ~/.bash_profile"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create Security Group"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash\n",
    "\n",
    "source ~/.bash_profile\n",
    "\n",
    "export SEC_GROUP_ID=$(aws ec2 create-security-group --group-name eks-fsx-security-group --vpc-id ${VPC_ID} --description \"FSx for Lustre Security Group\" --query \"GroupId\" --output text)\n",
    "echo \"export SEC_GROUP_ID=${SEC_GROUP_ID}\" | tee -a ~/.bash_profile"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Add an ingress rule that opens up port 988 from the 192.168.0.0/16 CIDR range"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash\n",
    "\n",
    "source ~/.bash_profile\n",
    "\n",
    "aws ec2 authorize-security-group-ingress --group-id ${SEC_GROUP_ID} --protocol tcp --port 988 --cidr 192.168.0.0/16"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Update the environment variables in the `storageclass.yaml` file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pygmentize fsx/storageclass.yaml"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Create FSX Storage Class"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#!kubectl delete -f fsx/storageclass.yaml"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!kubectl create -f fsx/storageclass.yaml"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!kubectl get sc"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Create Claim"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#!curl -o claim.yaml https://raw.githubusercontent.com/kubernetes-sigs/aws-fsx-csi-driver/master/examples/kubernetes/dynamic_provisioning_s3/specs/claim.yaml"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pygmentize fsx/claim.yaml"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#!kubectl delete -f fsx/claim.yaml"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!kubectl apply -f fsx/claim.yaml"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!kubectl get pvc fsx-claim"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!kubectl describe pvc fsx-claim"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## _Wait for status == Bound_"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Update FSX to `autoImportPolicy: NEW_CHANGED`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fsx = boto3.client(\"fsx\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "response = fsx.describe_file_systems()\n",
    "fsx_id = response[\"FileSystems\"][0][\"FileSystemId\"]\n",
    "print(fsx_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "response = fsx.update_file_system(FileSystemId=fsx_id, LustreConfiguration={\"AutoImportPolicy\": \"NEW_CHANGED\"})\n",
    "print(response)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "conda_python3",
   "language": "python",
   "name": "conda_python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
